import pandas as pd
import requests
import time
import os

def search_pubchem(compound_name):
    """
    Search PubChem for a compound by name and return the best match CID.
    """
    base_url = "https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name"
    search_url = f"{base_url}/{compound_name}/cids/JSON"
    
    try:
        response = requests.get(search_url)
        if response.status_code == 200:
            data = response.json()
            if 'IdentifierList' in data and 'CID' in data['IdentifierList']:
                # Return the first (best) match
                cid = data['IdentifierList']['CID'][0]
                print(f"Found CID for {compound_name}: {cid}")
                return cid
        print(f"No CID found for {compound_name}")
        return None
    except Exception as e:
        print(f"Error searching for '{compound_name}': {str(e)}")
        return None

def get_smiles_api(cid):
    """
    Get SMILES directly from PubChem PUG REST API.
    Much more reliable than web scraping.
    """
    if not cid:
        return None
    
    # Use PubChem's PUG REST API to get SMILES
    property_url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{cid}/property/CanonicalSMILES,IsomericSMILES/JSON"
    
    try:
        response = requests.get(property_url)
        if response.status_code == 200:
            data = response.json()
            if 'PropertyTable' in data and 'Properties' in data['PropertyTable']:
                properties = data['PropertyTable']['Properties'][0]
                
                # Prefer IsomericSMILES if available
                if 'IsomericSMILES' in properties and properties['IsomericSMILES']:
                    print(f"Found IsomericSMILES for CID {cid}: {properties['IsomericSMILES']}")
                    return properties['IsomericSMILES']
                
                # Fall back to CanonicalSMILES
                if 'CanonicalSMILES' in properties and properties['CanonicalSMILES']:
                    print(f"Found CanonicalSMILES for CID {cid}: {properties['CanonicalSMILES']}")
                    return properties['CanonicalSMILES']
            
            print(f"No SMILES found in API response for CID {cid}")
        else:
            print(f"API request failed with status code: {response.status_code}")
    except Exception as e:
        print(f"Error getting SMILES via API for CID {cid}: {str(e)}")
    
    return None

def get_molecular_formula_api(cid):
    """
    Get molecular formula using PubChem API.
    """
    if not cid:
        return None
    
    formula_url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{cid}/property/MolecularFormula/JSON"
    
    try:
        response = requests.get(formula_url)
        if response.status_code == 200:
            data = response.json()
            if 'PropertyTable' in data and 'Properties' in data['PropertyTable']:
                formula = data['PropertyTable']['Properties'][0]['MolecularFormula']
                print(f"Found molecular formula for CID {cid}: {formula}")
                return formula
        print(f"No molecular formula found for CID {cid}")
    except Exception as e:
        print(f"Error getting formula for CID {cid}: {str(e)}")
    
    return None

def process_excel_file(input_file, output_file=None):
    """
    Process the Excel file to get CID, MF, and SMILES using PubChem API.
    """
    if not output_file:
        base, ext = os.path.splitext(input_file)
        output_file = f"{base}_updated{ext}"
    
    try:
        # Read the Excel file
        print(f"Reading Excel file: {input_file}")
        df = pd.read_excel(input_file)
        print(f"Found {len(df)} rows in the Excel file")
        
        # Ensure required column exists
        if "Name_EN" not in df.columns:
            print("Error: 'Name_EN' column not found in the Excel file")
            return False
        
        # Ensure output columns exist
        output_columns = ["CID", "MF", "SMILES"]
        for col in output_columns:
            if col not in df.columns:
                df[col] = None
        
        # Process each row
        total_rows = len(df)
        for idx, row in df.iterrows():
            compound_name = row["Name_EN"]
            if pd.isna(compound_name) or not compound_name:
                continue
            
            print(f"Processing {idx+1}/{total_rows}: {compound_name}")
            
            # Skip if already processed
            if not pd.isna(row["CID"]) and not pd.isna(row["MF"]) and not pd.isna(row["SMILES"]):
                print(f"  Already processed {compound_name}, skipping...")
                continue
            
            # Step 1: Get CID
            cid = None
            if pd.isna(row["CID"]):
                cid = search_pubchem(compound_name)
                if cid:
                    df.at[idx, "CID"] = cid
                else:
                    print(f"  No CID found for {compound_name}, skipping...")
                    continue
            else:
                cid = int(row["CID"])
            
            # Step 2: Get Molecular Formula
            if pd.isna(row["MF"]):
                formula = get_molecular_formula_api(cid)
                if formula:
                    df.at[idx, "MF"] = formula
            
            # Step 3: Get SMILES
            if pd.isna(row["SMILES"]):
                smiles = get_smiles_api(cid)
                if smiles:
                    df.at[idx, "SMILES"] = smiles
            
            # Save regularly
            if (idx + 1) % 5 == 0 or idx == total_rows - 1:
                df.to_excel(output_file, index=False)
                print(f"Progress saved to {output_file}")
            
            # Delay to avoid overloading the API
            time.sleep(1)
        
        # Final save
        df.to_excel(output_file, index=False)
        print(f"Data processing complete. Updated file saved as: {output_file}")
        return True
        
    except Exception as e:
        print(f"Error processing Excel file: {str(e)}")
        return False

if __name__ == "__main__":
    # Specify your input file path
    input_file = "示例成分.xlsx"  # Change to your file path
    
    # Test with a specific compound
    test_compound = "Ferulic acid"
    print(f"Testing with compound: {test_compound}")
    
    # Step 1: Get CID
    cid = search_pubchem(test_compound)
    if cid:
        # Step 2: Get molecular formula
        formula = get_molecular_formula_api(cid)
        print(f"Molecular Formula: {formula}")
        
        # Step 3: Get SMILES
        smiles = get_smiles_api(cid)
        print(f"SMILES: {smiles}")
    
    # Uncomment to process the entire file
    process_excel_file(input_file)